#include "gtest/gtest.h"
#include "log.hpp"
#include "nn_detector.hpp"
#include "nn_yolo.hpp"
#include "opencv2/opencv.hpp"

// TEST(TestNN, TestYolo) {
//   component::Logger::SetLogger(
//       "logs/yolo.log", component::Logger::FMT::kFMT_TEST,
//       spdlog::level::debug);
//   TRTYoloBuilder builder("../../../data/nn/yolov5s.wts",
//                          "../../../data/yolo.engine");
//   TrtDetector detector;
//   detector.LoadEngine(builder.GetEngine());

//   for (unsigned int i = 1; i < 2; i++) {
//     std::string str = cv::format("../../../image/test/%d.bmp", i);
//     SPDLOG_WARN("{}", str);

//     detector.TestInfer(str);
//   }
// }

// TEST(TestNN, TestYoloSave) {
//   component::Logger::SetLogger(
//       "logs/yolo.log", component::Logger::FMT::kFMT_TEST,
//       spdlog::level::debug);
//   TRTYoloBuilder builder("../../../data/nn/yolov5s.wts",
//                       "../../../data/yolo.engine");
//   builder.SaveEngine();
// }

TEST(TestYolo, TestYoloWrite) {
  component::Logger::SetLogger("logs/yolo.log",
                               component::Logger::FMT::kFMT_DEFAULT);
  SPDLOG_WARN("Start");
  TRTYoloBuilder builder("", "../../../data/nn/yolo_mine_3090_mv.engine");
  builder.SetFilePath("../../../data/best_s.onnx");
  SPDLOG_WARN("Onnx");

  builder.CreateEngine(CreateMethod::kONNX);
  SPDLOG_WARN("CreateEngine");

  builder.SaveEngine();
  SPDLOG_WARN("SaveEngine");

  TRTDetector detector;
  detector.LoadEngine(builder.GetEngine());
  detector.Init();

  for (unsigned int i = 1; i < 28; i++) {
    std::string str = cv::format("../../../image/test/%d.bmp", i);
    SPDLOG_WARN("{}", str);

    detector.TestInfer(str);
  }
}

// TEST(TestNN, TestYoloRead) {
//   component::Logger::SetLogger(
//       "logs/yolo.log", component::Logger::FMT::kFMT_TEST,
//       spdlog::level::debug);

//   TRTDetector detector;
//   detector.SetEnginePath("../../../data/yolo.engine");
//   detector.Init();

//   for (unsigned int i = 1; i < 28; i++) {
//     std::string str = cv::format("../../../image/test/%d.bmp", i);
//     SPDLOG_WARN("{}", str);

//     detector.TestInfer(str);
//   }
// }